import os
from datasets import load_dataset, concatenate_datasets, load_from_disk

from prompt import get_industry_trans_func
from utils import glm4_vllm, load_obj
from setting import StaticValues
from parse import reason_label_parse

call_vllm = None


def run_llm_binary_cls(name, dataset, call_vllm, over_write=False):
    """
    从本地的大模型推理的pkl文件中
    大语言模型，针对少量数据做二分类

    * over_write: 是否覆盖本地已有的推理结果
    """
    sv = StaticValues(name)
    logger = sv.logger

    # 导入关键词分类的表
    # dataset = load_dataset("csv", data_files=sv.KW_CSV, split="train")
    attr_name = "is_industry_prompt"
    # 看文件是否存在
    if over_write or not os.path.exists(sv.BINARY_VLLM_OBJ):
        new_dataset = dataset.map(
            get_industry_trans_func(attr_name, sv.binary_cls_prompt),
        )
        prompts = []
        for item in new_dataset:
            prompts.append([{"role": "user", "content": item[attr_name]}])
            # 封装一下，如果已经数据，自动导入
        call_vllm(prompts, output_dir=sv.BINARY_VLLM_OBJ)
    else:
        logger.info(f"{sv.BINARY_VLLM_OBJ} exists, skip vllm")

    # 导入LLM的推理结果
    vllm_pred = load_obj(sv.BINARY_VLLM_OBJ)
    yes_cnt = 0
    no_cnt = 0

    reasons = []
    labels = []

    for idx, pred in enumerate(vllm_pred):
        pred_text = pred.outputs[0].text

        # 从大模型的预测文本中，解析出label
        reason, label = reason_label_parse(pred_text)
        reasons.append(reason)
        if label == "是":
            yes_cnt += 1
            labels.append(1)
        elif label == "否":
            no_cnt += 1
            labels.append(0)
        else:
            # 无法判断的数据，会给负样本数据引入错误的数据
            labels.append(0)

    if yes_cnt + no_cnt == len(vllm_pred):
        logger.info(
            f"在通过关键词随机筛选出的{len(vllm_pred)}个企业样本中，通过大模型分类，其中{yes_cnt}个是{sv.NAME}企业，{no_cnt}个不是{sv.NAME}企业。"
        )
    else:
        logger.info(
            f"在通过关键词随机筛选出的{len(vllm_pred)}个企业样本中，通过大模型分类，其中{yes_cnt}个是{sv.NAME}企业，{no_cnt}个不是{sv.NAME}企业。还有{len(vllm_pred) - yes_cnt - no_cnt}个样本无法确定。"
        )

    # 导出excel
    tmp_dataset = dataset.add_column("reason", reasons)
    tmp_dataset = tmp_dataset.add_column("label", labels)
    return tmp_dataset


def build_binary_dataset(name, dataset) -> object:
    """
    * 根据大模型分类结果，构建正负样本1:1的分类数据集
    * 导出正负样本，供用户修改二分类提示词
    *

    >>>

        dataset = load_from_disk(folder)["train"]
        # 选出所有
        build_binary_dataset(dataset, "test")

    """
    sv = StaticValues(name=name)

    # dataset = load_from_disk(folder)["train"]
    # 选出所有 label = 1的数据
    true_dataset = dataset.filter(lambda x: x["label"] == 1)
    true_dataset = true_dataset.map(
        get_industry_trans_func("industry_info", "{industry_info}")
    )

    # 正样本数据导出，仅供人查看，并辅助人工调整提示词
    true_dataset.to_csv(os.path.join(sv.human_folder, "企业二分类正样本.csv"))

    # 选出所有 label = -1的数据
    false_dataset = dataset.filter(lambda x: x["label"] == 0)

    # 由于负样本数据过多，只输出其中5000条数据。仅供人查看，并辅助人工调整提示词
    if len(false_dataset) > 5000:
        export_false_dataset = false_dataset.train_test_split(train_size=5000)["train"]
    else:
        export_false_dataset = false_dataset

    export_false_dataset.map(
        get_industry_trans_func("industry_info", "{industry_info}")
    ).to_csv(os.path.join(sv.human_folder, "企业二分类负样本.csv"))

    # 负样本数据，先 export，再变化 false_dataset
    # 若负样本的数量大于正样本的数量，切分负样本，让正负样本 1:1；
    # 若负样本数量小于正样本数量，无需切分
    if len(false_dataset) > len(true_dataset):
        false_dataset = false_dataset.train_test_split(train_size=len(true_dataset))["train"]

    # TODO： 没有考虑负样本小于正样本的情况，虽然一般这种情况也不会发生
    false_dataset = false_dataset.map(
        get_industry_trans_func("industry_info", "{industry_info}")
    )

    from datasets import concatenate_datasets
    bert_dataset = concatenate_datasets([true_dataset, false_dataset])
    bert_dataset = bert_dataset.shuffle(seed=42)

    return bert_dataset


def merge_binary_dataset(name):
    """
        合并大模型预测的数据和bert预测的数据，只留下正样本，负样本丢弃
    """
    sv = StaticValues(name=name)
    train_dataset = load_from_disk(sv.industry_dataset)["train"]
    binary_pred_csv = os.path.join(sv.home_folder, "bert_binary_pred.csv")

    if os.path.exists(binary_pred_csv):
        test_dataset = load_dataset("csv", data_files=binary_pred_csv, split="train")
        test_dataset = test_dataset.rename_column("binary_label", "label")

        dataset = concatenate_datasets([train_dataset, test_dataset])
    else:
        print(sv.chinese_name, "二分类全部由LLM推理")
        dataset = train_dataset
    dataset = dataset.filter(lambda x: x["label"] == 1)
    return dataset


if __name__ == "__main__":
    # run_llm_binary_cls("hydrogen")
    folder = "/home/jie/gitee/pku_industry/general/output/biomedical/industry_dataset"
    from datasets import load_from_disk

    dataset = load_from_disk(folder)["train"]
    # 选出所有
    tmp = build_binary_dataset(dataset, "test")
    print(tmp)
